'''
    Types
'''

try:
	import numpy as Numeric
	def sumDot( a,b ):
		return Numeric.dot (a, b)
except ImportError, err:
	try: 
		import Numeric
		def sumDot( a,b ):
			return sum (Numeric.dot (a, b) )
	except ImportError, err:
		print "Requires the numpy or Numeric extension"
		import sys
		sys.exit()
import copy
import math

###
### Base Types
###
class NumericObj:
    ''' Behave as Numeric.ndarray '''
    def __str__(self):
        return self._data
        
    def __getitem__(self, key):
        return self._data.__getitem__(key)
        
    def __setitem__(self, key, value):
        return self._data.__setitem__(key, value)
        
    def __call__(self):
        return self._data
        
    def copy(self):
        return copy.copy(self)

        
###
### Matrices
###

class Matrix(NumericObj):
    def transpose(self):
        self._data = Numeric.transpose(self._data)
    
    def T(self):
        return self.__class__(self._data, transpose = True)
    
class Matrix4(Matrix):
    ''' 4x4 Float matrix. Reference by row, then column '''
    def __init__(self, data = None, transpose = False):
        if data is None:
            self._data = Numeric.identity(4, 'f')
        else:
            self._data = Numeric.array(data)
            if transpose:
                self.transpose()
    
    def svd(self):
        ''' returns scale of Singular Value Decomposition '''
        # return Numeric.linalg.svd(self._data, compute_uv = False) - does not work...
        (X,Y,Z),data = range(3),self._data
        return math.sqrt ( 
            ( (data [X][X] * data [X][X]) + (data [X][Y] * data [X][Y]) + (data [X][Z] * data [X][Z]) +
            (data [Y][X] * data [Y][X]) + (data [Y][Y] * data [Y][Y]) + (data [Y][Z] * data [Y][Z]) +
            (data [Z][X] * data [Z][X]) + (data [Z][Y] * data [Z][Y]) + (data [Z][Z] * data [Z][Z]) ) / 3.0 )
            
            
    def embed(self, matrix):
        '''  
            Modifies the matrix in-place by replacing its upper 3x3 portion 
            with the passed in 3x3 matrix
        '''
        assert matrix.__class__ == Matrix3
        self._data[0:3,0:3] = matrix._data
        
    def getUpper(self):
        ''' Returns the upper 3x3 matrix '''
        return self._data[0:3, 0:3]
        
    def setRotationFromMatrix(self,matrix):
        ''' 
            Sets the rotational component (upper 3x3) of this matrix to the matrix
            values in the T precision Matrix3d argument; the other elements of
            this matrix are unchanged; a singular value decomposition is performed
            on this object's upper 3x3 matrix to factor out the scale, then this
            object's upper 3x3 matrix components are replaced by the passed rotation
            components, and then the scale is reapplied to the rotational
            components.
        '''
        scale = self.svd()
        self.embed(matrix)
        self.scale(scale)
        
        
    def scale(self, scale):
        ''' scale the matrix up by a scalar '''
        self._data *= scale
        
    def __mul__(self, vec4):
        ''' currently only supports dot product with a Vector4 or a tuple (however, length is not tested) '''
        if isinstance(vec4, Vector4):
            return Numeric.dot(self._data, vec4._data)
        elif type(vec4) is tuple:
            return Numeric.dot(self._data, vec4)
        else:
            raise Exception('Matrix multiplication with this type is not supported')
        
class Matrix3(Matrix):
    ''' 3x3 Float matrix. Reference by row, then column '''
    def __init__(self, data = None, transpose = False):
        if data is None:
            self._data = Numeric.identity(3, 'f')
        else:
            self._data = Numeric.array(data)
            self._data.shape = (3,3)
            if transpose:
                self.transpose()
        
    @staticmethod
    def Multiply(a, b):
        ''' Multiplies two 3x3 matrices '''
        assert a.__class__ == Matrix3 and b.__class__ == Matrix3, "Arguments of type %s, %s" % (a.__class__.__name__, b.__class__.__name__)
        result = Matrix3()
        result._data = sumDot(a._data, b._data)
        return result
        
    @staticmethod 
    def FromQuaternion(quaternion):
        ''' Returns a 3x3 Rotation matrix from H quaternion) '''
        X,Y,Z,W = range(4)
        q1 = quaternion

        result = Matrix3()
        n = q1.dot(q1)
        s = 0.0
        if (n > 0.0):
            s = 2.0 / n
        xs = q1 [X] * s;  ys = q1 [Y] * s;  zs = q1 [Z] * s
        wx = q1 [W] * xs; wy = q1 [W] * ys; wz = q1 [W] * zs
        xx = q1 [X] * xs; xy = q1 [X] * ys; xz = q1 [X] * zs
        yy = q1 [Y] * ys; yz = q1 [Y] * zs; zz = q1 [Z] * zs
        # This math all comes about by way of algebra, complex math, and trig identities.
        # See Lengyel pages 88-92
        result [X][X] = 1.0 - (yy + zz);	result [Y][X] = xy - wz; 			result [Z][X] = xz + wy;
        result [X][Y] =       xy + wz; 		result [Y][Y] = 1.0 - (xx + zz);	result [Z][Y] = yz - wx;
        result [X][Z] =       xz - wy; 		result [Y][Z] = yz + wx;          	result [Z][Z] = 1.0 - (xx + yy)

        return result
    
    @staticmethod
    def RotationZ(angle):
        ''' angle is in radians '''
        s, c = math.sin(angle), math.cos(angle)
        return Matrix3([c, -s, 0., s, c, 0., 0., 0., 1.])
    
    @staticmethod
    def RotationY(angle):
        ''' angle is in radians '''
        s, c = math.sin(angle), math.cos(angle)
        return Matrix3([c, 0., s, 0., 1., 0., -s, 0., c])
        
    @staticmethod
    def RotationX(angle):
        ''' angle is in radians '''
        s, c = math.sin(angle), math.cos(angle)
        return Matrix3([1., 0., 0., 0., c, -s, 0., s, c])
        
    def __mul__(self, other):
        return Matrix3.Multiply(self, other)
# Synonyms
Matrix3f = Matrix3
Matrix4f = Matrix4
        
###
### Vectors
###

class Vector(object):
    ''' Vector base class '''
    def __repr__(self):
        ''' display as underlying data '''
        return repr(self._data)
    
    def __str__(self):
        ''' display as underlying data (str) '''
        return str(self._data)
        
    def __length__(self):
        ''' returns the number of rows, NOT LENGTH! '''
        return len(self._data)
        
    def dot(self, other):
        ''' Dot product '''
        assert isinstance(other, self.__class__) and isinstance(self, other.__class__)
        return Numeric.dot(self._data,other._data)
        
    def length(self):
        ''' returns vector length '''
        return math.sqrt(self.dot(self))

    def scale(self, scalar):
        ''' in-place scalar multiplication '''
        self._data *= scalar
    
    def normalize(self):
        self.scale(1.0/self.length())
    
    def __add__(self, other):
        ''' addition '''
        assert isinstance(other, self.__class__) and isinstance(self, other.__class__)
        return self.__class__(*(self._data + other._data))
        
    def __sub__(self, other):
        assert isinstance(other, self.__class__) and isinstance(self, other.__class__)
        return self.__class__(*(self._data - other._data))
    
    def __mul__(self, other):
        # type unsafety
        if type(other)==float:
            return self.__class__(*Numeric.dot(self._data, other))
        else:
            return self.__class__(*Numeric.dot(self._data, other._data))
        
    
    def __iadd__(self, other):
        assert isinstance(other, self.__class__) and isinstance(self, other.__class__)
        self._data += other._data
        return self
    
    def __isub__(self, other):
        assert isinstance(other, self.__class__) and isinstance(self, other.__class__)
        self._data -= other._data
        return self
        
    @staticmethod
    def attr_fn(i):
        ''' Helper function to generate setters/getters for vector attributes '''
        def fget(self):
            return self._data[i]
        def fset(self, value):
            self._data[i] = value
        return fget, fset


class Vector2(NumericObj, Vector):
    ''' 2-row vector '''
    def __init__(self, x = 0.0, y = 0.0):    
        pt = Numeric.zeros(2, 'f')
        pt[0] = x
        pt[1] = y
        self._data = pt
    
    x = property(*Vector.attr_fn(0))
    y = property(*Vector.attr_fn(1))
    
    u = property(*Vector.attr_fn(0))
    v = property(*Vector.attr_fn(1))
    
    s = property(*Vector.attr_fn(0))
    t = property(*Vector.attr_fn(1))

class Vector3(Vector2):
    ''' 3-row vector '''
    def __init__(self, x = 0.0, y = 0.0, z = 0.0):
        pt = Numeric.zeros(3, 'f')
        pt[0] = x
        pt[1] = y
        pt[2] = z
        self._data = pt
        
    z = property(*Vector.attr_fn(2))
    
    r = property(*Vector.attr_fn(0))
    g = property(*Vector.attr_fn(1))
    b = property(*Vector.attr_fn(2))
    
    def setTo(self, x, y, z):
        self.x, self.y, self.z = x, y, z
    
    @staticmethod
    def Cross(u, v):
        assert u.__class__ == Vector3 and v.__class__ == Vector3
        # Cross product of two 3f vectors
        X,Y,Z = range(3)
        cross = Vector3()
        cross[X] = (u[Y] * v[Z]) - (u[Z] * v[Y])
        cross[Y] = (u[Z] * v[X]) - (u[X] * v[Z])
        cross[Z] = (u[X] * v[Y]) - (u[Y] * v[X])
        return cross
        
    def __len__(self):
        return 3
        
class Vector4(Vector3):
    ''' 4-row vector '''
    def __init__(self, x = 0.0, y = 0.0, z = 0.0, a = 0.0):
        pt = Numeric.zeros(4, 'f')
        pt[0] = x
        pt[1] = y
        pt[2] = z
        pt[3] = a
        self._data = pt

    w = property(*Vector.attr_fn(3))
    
    a = property(*Vector.attr_fn(3))
    
    def __len__(self):
        return 4

# Synonyms
Point2f         = Vector2
Vector2f        = Vector2
Point3f         = Vector3
Vector3f        = Vector3
Quaternion4f    = Vector4